from pprint import pprint
import paramiko
from io import StringIO
import threading
from threading import Thread
import traceback
import socket
import json

from paramiko.ssh_exception import NoValidConnectionsError
from paramiko.ssh_exception import AuthenticationException
from paramiko.ssh_exception import SSHException
zmodemszstart = b'rz\r**\x18B00000000000000\r\x8a'
zmodemszend = b'**\x18B0800000000022d\r\x8a'
zmodemrzstart = b'rz waiting to receive.**\x18B0100000023be50\r\x8a'
zmodemrzend = b'**\x18B0800000000022d\r\x8a'
zmodemcancel = b'\x18\x18\x18\x18\x18\x08\x08\x08\x08\x08'
from utils.custom_log import log_start
logger = log_start('ssh')

class SSH:
    def __init__(self, websocker, message):
        self.websocker = websocker
        self.message = message
        self.cmd = ''
        self.res = ''
        self.zmodem = False
        self.zmodemOO = False
    
    # term 可以使用 ansi, linux, vt100, xterm, dumb,除了 dumb外其他都有颜色显示
    def connect(self, host, user, password=None, ssh_key=None, port=22, timeout=None,width=None,height=None):
        
        try:
            # 实例化SSHClient
            logger.info('实例化SSHClient')
            ssh_client = paramiko.SSHClient()
            # 当远程服务器没有本地主机的密钥时自动添加到本地,这样不用在建立连接的时候输入yes或no进行确认
            ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            # 用key进行认证
            if ssh_key:
                private = StringIO(ssh_key)
                key = paramiko.RSAKey.from_private_key(private)
                ssh_client.connect(hostname=host, port=port, username=user, pkey=key)
                logger.info(f'sshkey认证登录:ssh {user}@{host} -p {port}')
                # look_for_keys: 是否在~/.ssh中搜索私钥文件
                # allow_agent: 是否允许连接到ssh代理
            else:
                ssh_client.connect(hostname=host, port=port, username=user, password=password,
                timeout=timeout)
                logger.info(f'密码认证登录:ssh {user}@{host} -p {port}')


            # 打开ssh通道,建立长连接
            logger.info("transport.open_session() 打开ssh通道,建立长连接")
            transport = ssh_client.get_transport()
            self.channel = transport.open_session()
            # 获取ssh通道,并设置term和终端大小
            self.channel.get_pty(term='xterm', width=width, height=height)
            # 激活终端SSH通道
            self.channel.invoke_shell()

            # 一开始展示Linux欢迎相关内容,后面不进入此方法
            for i in range(2):
                recv = self.channel.recv(1024).decode('utf-8')
                self.message['status'] = 0
                self.message['message'] = recv
                message = json.dumps(self.message)
                self.websocker.send(message)
                self.res += recv
                logger.info(self.res)

            # 创建3个线程将服务器返回的数据发送到django websocket（1个线程都可以）
            Thread(target=self.websocket_to_django).start()
            Thread(target=self.websocket_to_django).start()
            Thread(target=self.websocket_to_django).start()

        except NoValidConnectionsError as connect:
            msg = {'status': 1, 'message': "无法建立SSH连接:%s,请重试！" % connect}
            logger.error(msg)
            self.websocker.send(json.dumps(msg))
            self.websocker.close(3001)
           

        except AuthenticationException as auth:
            msg = {'status': 1, 'message': '验证错误: %s,检查用户密码!!! ' % auth}
            logger.error(msg)
            self.websocker.send(json.dumps(msg))
            self.websocker.close(3001)
           

        except SSHException as ssh:
            msg = {'status': 1, 'message': "SSH连接失败: %s,请检查重试" % ssh}
            logger.error(msg)
            self.websocker.send(json.dumps(msg))
            self.websocker.close(3001)
           

        except Exception as  e:
            msg = {'status': 3, 'message': "%s"  % e}
            logger.error(msg)
            self.websocker.close(3001)
  

    def resize_pty(self, cols, rows):
        logger.info(f"resize_pty动态改变Terminal窗口的大小;cols: [{cols}], rows: [{rows}]")
        self.channel.resize_pty(width=cols, height=rows)

    def django_to_ssh(self, data):
        try:
            logger.info(f'SSH channel send::,{data}')
            self.channel.send(data)
            if data == '\r':
                data = '\n'
            self.cmd += data
        except Exception as  e:
            self.close('连接SSH异常,关闭连接,%s' % e)

    def django_bytes_to_ssh(self, data):
        try:
            self.channel.send(data)
        except Exception  as  e:
            self.close('连接SSH异常,关闭连接,%s' % e)

    def websocket_to_django(self):
        try:
            while not self.channel.exit_status_ready():
                if self.zmodemOO:
                    # 文件开始下载
                    self.zmodemOO = False
                    data = self.channel.recv(2)
                    if not len(data):
                        return
                    # 文件下载结束
                    if data == b'OO':
                        self.websocker.send(bytes_data=data)
                        continue
                    else:
                        data = data + self.channel.recv(40960)
                else:
                    data = self.channel.recv(40960)
                    if not len(data):
                        return

                if self.zmodem:
                    if zmodemszend in data or zmodemrzend in data:
                        self.zmodem = False
                        if zmodemszend in data:
                            self.zmodemOO = True
                    if zmodemcancel in data:
                        self.zmodem = False
                    self.websocker.send(bytes_data=data)
                else:
                    if zmodemszstart in data or zmodemrzstart in data:
                        self.zmodem = True
                        self.websocker.send(bytes_data=data)
                    else:
                        # SSH返回的数据需要转码为utf-8,否则json序列化会失败
                        data = data.decode('utf-8')
                        self.message['status'] = 0
                        self.message['message'] = data
                        self.res += data
                        message = json.dumps(self.message)
                        logger.info(f"websockerSend 将返回的$message输出:{message}")
                        self.websocker.send(message)
        except Exception as  e:
            self.close(f'发送到WSS异常,关闭连接: {e} data: {data}')

    def close(self,message='关闭连接'):
        self.message['status'] = 1
        self.message['message'] =  message
        message = json.dumps(self.message)
        logger.info(f"关闭连接:{message}")
        self.websocker.send(message)
        self.websocker.close()
        self.channel.close()

    def shell(self, data):
        # 创建线程的方式发送数据到ssh,每次发送都是一个字符,可以不用线程
        # 直接调用函数性能更好
        # self.django_to_ssh(data)
        Thread(target=self.django_to_ssh, args=(data,)).start()
        # Thread(target=self.websocket_to_django).start()
        # Thread(target=self.websocket_to_django).start()
        # Thread(target=self.websocket_to_django).start()